import mlflow
import pandas as pd

import mlflow
import pandas as pd

def generate_recommendations_table(experiment_ids, prefix_note="sizes_acts", group_type="sim"):
    all_rows = []

    for exp_id in experiment_ids:
        runs = mlflow.search_runs(
            experiment_ids=[exp_id],
            output_format="list"
        )
        for run in runs:
            if prefix_note not in run.data.params.get("note") or run.data.params.get("group_type") != group_type:
                continue

            dim = int(run.data.params.get("embedding_dim", 0))
            dataset = run.data.params.get("dataset", f"Exp-{exp_id}")
            topk = int(run.data.params.get("top_k", 0))
            activation = run.data.params.get("topk_inference", "False")
            aggregation_function = run.data.params.get("SAE_fusion_strategy", "none")
            

            row_key = (dim, topk, aggregation_function, activation)
            metrics = {
                (dataset, "G/mean"): run.data.metrics.get("CommonItemsNDCG20/median"),
                (dataset, "U/mean"): run.data.metrics.get("NDCG20/mean"),
                (dataset, "U/min"): run.data.metrics.get("NDCG20/min"),
                (dataset, "Pop"): run.data.metrics.get("Popularity/mean"),
            }

            all_rows.append((row_key, metrics))

    # Build DataFrame from records
    records = {}
    for key, metrics in all_rows:
        if key not in records:
            records[key] = {}
        records[key].update(metrics)

    df = pd.DataFrame.from_dict(records, orient="index")
    df.index.names = ["Dimensions", "TopK", "Aggregation", "Activation"]

    # Sort and reindex columns by dataset then metric
    df = df.sort_index(axis=1, level=[0, 1]).sort_values(
        by=["Dimensions", "TopK", "Aggregation", "Activation"]
    )

    return df.reset_index()

def highlight_top3_dark_to_light(s):
    # Colors from dark to light
    colors = ['mediumseagreen', 'lightgreen']
    
    # Get sorted unique values in descending order
    top_values = s.nlargest(2).unique()
    
    # Assign background color depending on rank
    styles = ['' for _ in s]
    for rank, value in enumerate(top_values):
        styles = [
            f'background-color: {colors[rank]}' if v == value and styles[i] == '' else styles[i]
            for i, v in enumerate(s)
        ]
    return styles

def highlight_bottom3_dark_to_light(s):
    # Colors from dark to light
    colors = ['mediumblue', 'lightblue', 'paleturquoise']
    
    # Get sorted unique values in ascending order
    bottom_values = s.nsmallest(3).unique()
    
    # Assign background color depending on rank
    styles = ['' for _ in s]
    for rank, value in enumerate(bottom_values):
        styles = [
            f'background-color: {colors[rank]}' if v == value and styles[i] == '' else styles[i]
            for i, v in enumerate(s)
        ]
    return styles

experiments = ['523100174176986081', '333391697323445885']

# Select only the desired columns for aggregation


table = generate_recommendations_table(experiments, prefix_note="aggregations", group_type="sim")

row_indexes_selected = [
    ('average', 'True'),
    ('common_features', 'False'),
    ('max', 'True'),
    ('square_average', 'False'),
    ('topk', 'False'),
    ('wcom', 'True'),
]

table = table[
    table.set_index(['Aggregation', 'Activation']).index.isin(row_indexes_selected)
].reset_index()


selected_columns = []
for dataset in ["MovieLens", "LastFM1k"]:
    for metric in ["G/mean", "U/mean", "U/min", "Pop"]:
        selected_columns.append((dataset, metric))
# Group by Aggregation and Activation and aggregate

agg_table = (
    table
    .groupby(["Aggregation", "Activation"])[selected_columns]
    .agg(['mean'])
).round(3)

selected_columns = []
for dataset in ["MovieLens", "LastFM1k"]:
    for metric in ["G/mean", "U/mean", "U/min"]:
        selected_columns.append((dataset, metric, 'mean'))
    
std_selected_columns = []
for dataset in ["MovieLens", "LastFM1k"]:
    for metric in ["G/mean", "U/mean", "U/min"]:
        std_selected_columns.append((dataset, metric, 'std'))
    


agg_table.style.apply(highlight_top3_dark_to_light, subset=selected_columns)#.highlight_min((std_selected_columns))

experiments = ['523100174176986081', '333391697323445885']

# Select only the desired columns for aggregation


table = generate_recommendations_table(experiments, prefix_note="aggregations", group_type="random")

row_indexes_selected = [
    ('average', 'True'),
    ('common_features', 'False'),
    ('max', 'True'),
    ('square_average', 'False'),
    ('topk', 'False'),
    ('wcom', 'True'),
]

table = table[
    table.set_index(['Aggregation', 'Activation']).index.isin(row_indexes_selected)
].reset_index()


selected_columns = []
for dataset in ["MovieLens", "LastFM1k"]:
    for metric in ["G/mean", "U/mean", "U/min", "Pop"]:
        selected_columns.append((dataset, metric))
# Group by Aggregation and Activation and aggregate

agg_table = (
    table
    .groupby(["Aggregation", "Activation"])[selected_columns]
    .agg(['mean'])
).round(3)

selected_columns = []
for dataset in ["MovieLens", "LastFM1k"]:
    for metric in ["G/mean", "U/mean", "U/min"]:
        selected_columns.append((dataset, metric, 'mean'))
    
std_selected_columns = []
for dataset in ["MovieLens", "LastFM1k"]:
    for metric in ["G/mean", "U/mean", "U/min"]:
        std_selected_columns.append((dataset, metric, 'std'))
    


agg_table.style.apply(highlight_top3_dark_to_light, subset=selected_columns)#.highlight_min((std_selected_columns))

experiments = ['523100174176986081', '333391697323445885']

# Select only the desired columns for aggregation


table = generate_recommendations_table(experiments, prefix_note="aggregations", group_type="div")

row_indexes_selected = [
    ('average', 'True'),
    ('common_features', 'False'),
    ('max', 'True'),
    ('square_average', 'False'),
    ('topk', 'False'),
    ('wcom', 'True'),
]

table = table[
    table.set_index(['Aggregation', 'Activation']).index.isin(row_indexes_selected)
].reset_index()


selected_columns = []
for dataset in ["MovieLens", "LastFM1k"]:
    for metric in ["G/mean", "U/mean", "U/min", "Pop"]:
        selected_columns.append((dataset, metric))
# Group by Aggregation and Activation and aggregate

agg_table = (
    table
    .groupby(["Aggregation", "Activation"])[selected_columns]
    .agg(['mean'])
).round(3)

selected_columns = []
for dataset in ["MovieLens", "LastFM1k"]:
    for metric in ["G/mean", "U/mean", "U/min"]:
        selected_columns.append((dataset, metric, 'mean'))
    
std_selected_columns = []
for dataset in ["MovieLens", "LastFM1k"]:
    for metric in ["G/mean", "U/mean", "U/min"]:
        std_selected_columns.append((dataset, metric, 'std'))
    


agg_table.style.apply(highlight_top3_dark_to_light, subset=selected_columns)#.highlight_min((std_selected_columns))